#!/usr/bin/env python

import optparse
import sys

import tables

from charmicat import BadCommandLineError

def main():
    '''
    Usage: ypipi_kekpnl_search.py [opts] event_ids.txt file.count.h5
    
    Searches for run/event numbers listed in `event_ids.txt' (in format 
    RUN <any whitespace> EVENT, one pair per line) in `file.count.h5'.
    '''

    op = optparse.OptionParser(usage=main.__doc__)    
    
    op.add_option('-o', '--output', dest='output', default='ypipi_kekpnl_search.h5',
                  help='Put PyTables output in FILE', metavar='FILE')

    opts, args = op.parse_args()

    if len(args) < 2:
        msg = "Please supply a text file with run/event numbers and a PyTables file with a `count' table inside."
        raise BadCommandLineError(msg, -1)

    event_ids = []
    with open(args[0]) as fp:
        event_ids = [tuple(line.strip().split()) for line in fp if line.strip()]
        
    countdict = {}
    with tables.openFile(args[1]) as fp:
        with tables.openFile(opts.output, 'w') as fpout:
            nt = fp.root.nt
            ntout = fpout.createTable(fpout.root, 'nt', nt.coldescrs)

            for i, evid in enumerate(event_ids):
                sys.stdout.write("{0}/{1} queries complete\r".format(i, len(event_ids)))
                sys.stdout.flush()

                sel = '(runno == {}) & (eventno == {})'.format(*evid)

                countdict[evid] = nt.whereAppend(ntout, sel)
            
            ntout.flush()

    print "{0}/{0} queries complete\r".format(len(event_ids))

    keys = countdict.keys()
    keys.sort()

    for k in keys:
        print "{0:2d} {1:10d} {2:3d}".format(*map(int, [k[0], k[1], countdict[k]]))

    return 0


if __name__ == '__main__':
    try:
        import sys
        sys.exit(main())
    except BadCommandLineError, e:
        import os
        sys.stderr.write(str(e) + "\n")
        sys.stderr.write("See `{0} -h' for more.\n".format(os.path.basename(sys.argv[0])))
        sys.stderr.flush()
        sys.exit(e.retcode)
            
